library(readxl)
library(tidyverse)
library(swimplot)
library(ggrepel)
library(ggsci)

# %% Load data and processing ----
source("x28_help_funs.R")

path <- "D:/Work/TF/Ph1" # for working computer
datafile <- "XNW28012_SubjectList.xlsx"

# Load patients data
patient.info <- read_patient_info(path, datafile)

# %% Plotting functions ----

#' Draw waterfall plot
#'
#' @param plot.set: a data frame of selected patients
#' @param title: the title of the plot
#' @param path: the directory of the data file
#' @param datafile: the data file name
#' @return: an object of waterfall plot
ggwaterfall <- function(
  plot.set,
  title = "Waterfall",
  path = "D:/Work/TF/Ph1",
  datafile = "XNW28012_SubjectList.xlsx"
) {
  # RECIST data processing ---
  # Load RECIST data
  recist <- read_targetlesion_info(path, datafile)

  # Delete Not Evaluation Patients
  recist <- recist |>
    filter(Sum >= 0) # Del NA (Del patients without target lesions)

  # Baseline
  baseline <- recist |>
    filter(No.eva == 0)

  # Get the sum of target lesions of the best response
  best.resp <- recist |>
    filter(No.eva > 0) |>
    group_by(Patient) |>
    summarize(Min = min(Sum))

  # Calculate the Max change from baseline
  #   then join the patients info
  waterfall <- left_join(best.resp, baseline, by = "Patient") |>
    mutate(MaxChange = (Min - Sum) / Sum * 100) |>
    select(Patient, MaxChange) |>
    left_join(patient.info, by = "Patient")

  # Select Plot Set ----
  waterfall <- filter(waterfall, Patient %in% plot.set$Patient)

  # Draw Waterfall plot ----
  ## Tumor response with Measurable Disease
  wfplot <- ggplot(
    data = waterfall,
    aes(x = reorder(Patient, -MaxChange), y = MaxChange)
  ) +
    geom_col(aes(fill = Dose)) +
    geom_hline(yintercept = c(-30, 20), color = "grey50", linetype = "dashed") +
    labs(
      title = title,
      x = "",
      y = "% Change from Baseline in Sum of Diameters"
    ) +
    scale_fill_manual(values = DOSE.COLORS) +
    theme_bw()

  # # 添加BOR信息
  #   waterfall.bor.text <- waterfall |>
  #     filter(BOR != "SD")
  #   wfplot <- wfplot +
  #     # labs(title = tumor_set) +
  #     geom_text(data = waterfall.bor.text,
  #       aes(label = BOR),
  #       nudge_x = 0,
  #       nudge_y = 0,
  #       angle = 45,
  #       size = 2
  #     )+
  #     # scale_color_manual(values = RESP_COLORS) + #
  #     theme(axis.text.x = element_text(angle=90, size = 6)) # Show Patient ID

  # 对于某些增大<100%的情况，调整y轴范围，以保持对称
  max_change <- max(waterfall$MaxChange, na.rm = TRUE)
  if (max_change < 100) {
    wfplot <- wfplot +
      ylim(c(-100, 100))
  } else {
    wfplot <- wfplot +
      ylim(c(-100, max(waterfall$MaxChange, na.rm = TRUE) * 1.1))
  }

  # # Move legend to the right top of the figure
  # wfplot <- wfplot +
  #   theme(legend.position = c(0.975, 0.975),
  #         legend.justification = c(1, 1),
  #         legend.background = element_rect(fill = "transparent"))

  print("Waterfall plot")
  print(paste0(waterfall |> summarise(np = n_distinct(Patient)), " patients"))
  print(paste0(waterfall |> summarise(n = n()), " data points."))

  wfplot <- wfplot +
    theme(axis.text.x = element_text(angle = 90, hjust = 1, vjust = 0.5))
  
  return(wfplot)
}


#' Draw spider plot (baseline is 0)
#'
#' @param plot.set: a data frame of selected patients
#'
#' @param path: the directory of the data file
#' @param datafile: the data file name
#' @return: an object of spider plot
ggspider <- function(
  plot.set,
  title = "Spider plot",
  path = "D:/Work/TF/Ph1",
  datafile = "XNW28012_SubjectList.xlsx"
) {
  # ggspider()中基线的靶病灶尺寸归一化为0
  # 这是与ggspider2()的唯一区别

  # recist means all efficacy evaluations
  recist <- read_targetlesion_info(path, datafile)

  # 仅保留有效的评估患者
  recist <- recist |>
    filter(Sum >= 0) # Delete Not Evaluation Patients

  baseline <- recist |> filter(No.eva == 0)

  lesion <- recist |>
    filter(No.eva > 0) |>
    left_join(baseline, by = "Patient") |>
    mutate(
      Patient = Patient,
      # Transfer <duration> to number
      #   xxx.y -> Baseline
      Time = as.numeric(Date.eva.x - Date.eva.y, "weeks"),
      # Size是每次靶病灶径长和 与 基线径长和 的差值与基线的百分比
      #   图中基线是0
      Size = (Sum.x - Sum.y) / Sum.y * 100
    )

  # firstEva.df is the start points for Spider plot
  first.evaluate <- baseline |>
    mutate(Time = 0, Size = 1) |> # or Size=1
    select(Patient, Time, Size)

  # Data for plot ----
  spider <- first.evaluate |>
    bind_rows(lesion) |>
    left_join(patient.info, by = "Patient") |>
    select(Patient, Time, Tumor, Size, RECIST.x, BOR, Dose, isDoseEsca)

  max.duration <- max(spider$Time, na.rm = TRUE)

  # Data Set ----
  spider <- filter(spider, Patient %in% plot.set$Patient)

  # 删除只有baseline，没有后续肿评的受试者
  spider <- spider |>
    group_by(Patient) |>
    mutate(n = n()) |>
    filter(n > 1) |>
    ungroup()

  # Draw ----
  spplot <- ggplot(data = spider, aes(x = Time, y = Size)) +
    geom_point(aes(shape = RECIST.x, color = RECIST.x), size = 2, alpha = 0.9) +
    geom_line(aes(group = Patient, color = BOR), linewidth = 1, alpha = 0.7)
  # labs(title = tumor_set) +

  # Modify Y axis range (-100%, max size)
  max.size <- max(spider$Size, na.rm = TRUE) # 增大是是有可能>100%
  if (max.size < 100) {
    max.size <- 100
  } else {
    max.size <- max.size * 1.1
  }
  spplot <- spplot +
    ylim(c(-100, max.size))

  # 对于某些增大<2的情况，调整y轴范围，以保持对称
  # if (max(spider$Size, na.rm = TRUE) < 2) {
  #   spplot <- spplot +
  #     ylim(c(0, 2))
  # }

  # Modify the response shape (without NE)
  spplot <- spplot +
    theme_bw() +
    scale_shape_manual(
      name = "RECIST.x",
      values = RESP.SHAPES, # Shapes for CR, PR, SD, PD, NE
      breaks = RESP.LEVELS
    )

  spplot <- spplot +
    geom_hline(
      yintercept = c(20, -30), # PD and PR line
      color = "grey50",
      linetype = "dashed"
    ) +
    expand_limits(x = 0, y = 1) + # from (0, 1)
    labs(title = title, x = "Weeks", y = "Tumor Size to Baseline") +
    scale_x_continuous(
      breaks = seq(0, max.duration + 1, 6), # X axis by 6 weeks
      minor_breaks = seq(0, max.duration + 1, 3)
    ) +
    # 将shape和color图例合并，并设标题为空
    guides(shape = guide_legend(title = ""), color = guide_legend(title = ""))
  # 如果出现shape和color图例未能合并的情况，
  # 很可能是Patient (XNW28)表和TargetLesion表不一致

  # # Move legend to top right of the figure
  # spplot <- spplot +
  #   theme(legend.position = c(0.975, 0.8),
  #         legend.jusitification = c(1, 1),
  #         legend.background = element_rect(fill = "transparent"))

# 如果需要，可以加上受试者ID
# spplot <- spplot +
#     geom_text_repel(aes(label = Patient), alpha = 0.5)

  # Print summary
  print("Spider plot")
  print(paste0(spider |> summarise(np = n_distinct(Patient)), " patients"))
  print(paste0(spider |> summarise(n = n()), " data points."))

  return(spplot)
}


#' Draw spider plot (baseline is 1)
#'
#' @param plot.set: a data frame of selected patients
#'
#' @param path: the directory of the data file
#' @param datafile: the data file name
#' @return: an object of spider plot
ggspider2 <- function(
  plot.set,
  title = "Spider plot",
  path = "D:/Work/TF/Ph1",
  datafile = "XNW28012_SubjectList.xlsx"
) {
  # ggspider()中基线的靶病灶尺寸归一化为1
  # 这是与ggspider()的唯一区别

  # recist means all efficacy evaluations
  recist <- read_targetlesion_info(path, datafile)

  # 仅保留有效的评估患者
  recist <- recist |>
    filter(Sum >= 0) # Delete Not Evaluation Patients

  baseline <- recist |> filter(No.eva == 0)

  lesion <- recist |>
    filter(No.eva > 0) |>
    left_join(baseline, by = "Patient") |>
    mutate(
      Patient = Patient,
      # Transfer <duration> to number
      #   xxx.y -> Baseline
      Time = as.numeric(Date.eva.x - Date.eva.y, "weeks"),
      # Size是每次靶病灶径长和 与 基线径长和 的比值
      #  图中基线是1
      Size = (Sum.x / Sum.y),
    )

  # firstEva.df is the start points for Spider plot
  first.evaluate <- baseline |>
    mutate(Time = 0, Size = 1) |> # or Size=1
    select(Patient, Time, Size)

  # Data for plot ----
  spider <- first.evaluate |>
    bind_rows(lesion) |>
    left_join(patient.info, by = "Patient") |>
    select(Patient, Time, Tumor, Size, RECIST.x, BOR, Dose, isDoseEsca)

  max.duration <- max(spider$Time, na.rm = TRUE)

  # Data Set ----
  spider <- filter(spider, Patient %in% plot.set$Patient)

  # 删除只有baseline，没有后续肿评的受试者
  spider <- spider |>
    group_by(Patient) |>
    mutate(n = n()) |>
    filter(n > 1) |>
    ungroup()

  # Draw ----
  spplot <- ggplot(data = spider, aes(x = Time, y = Size)) +
    geom_point(aes(shape = RECIST.x, color = RECIST.x), size = 2, alpha = 0.9) +
    geom_line(aes(group = Patient, color = BOR), linewidth = 0.5, alpha = 0.7)
  # labs(title = tumor_set) +

  # 对于某些增大<2的情况，调整y轴范围，以保持对称
  if (max(spider$Size, na.rm = TRUE) < 2) {
    spplot <- spplot +
      ylim(c(0, 2))
  }

  # Modify the response shape (without NE)
  spplot <- spplot +
    theme_bw() +
    scale_shape_manual(
      name = "RECIST.x",
      values = RESP.SHAPES, # Shapes for CR, PR, SD, PD, NE
      breaks = RESP.LEVELS
    )

  spplot <- spplot +
    geom_hline(
      yintercept = c(1.2, 0.7), # PD and PR line
      color = "grey50",
      linetype = "dashed"
    ) +
    expand_limits(x = 0, y = 1) + # from (0, 1)
    labs(title = title, x = "Weeks", y = "Tumor Size to Baseline") +
    scale_x_continuous(
      breaks = seq(0, max.duration + 1, 6), # X axis by 6 weeks
      minor_breaks = seq(0, max.duration + 1, 3)
    ) +
    # 将shape和color图例合并，并设标题为空
    guides(shape = guide_legend(title = ""), color = guide_legend(title = ""))
  # 如果出现shape和color图例未能合并的情况，
  # 很可能是Patient (XNW28)表和TargetLesion表不一致

  # # Move legend to top right of the figure
  # spplot <- spplot +
  #   theme(legend.position = c(0.975, 0.8),
  #         legend.jusitification = c(1, 1),
  #         legend.background = element_rect(fill = "transparent"))

  # Print summary
  print("Spider plot")
  print(paste0(spider |> summarise(np = n_distinct(Patient)), " patients"))
  print(paste0(spider |> summarise(n = n()), " data points."))

  return(spplot)
}


#' Plot swimmer plot
#' @param plot.set: a data frame of patients information
#' @param title: plot title
#' @param path: the directory of the data file
#' @param datafile: the data file name
#' @return: an object of swimmer plot
ggswim <- function(
  plot.set,
  title = "Swimmer Plot",
  path = "D:/Work/TF/Ph1",
  datafile = "XNW28012_SubjectList.xlsx"
) {
  # Raw data (treatment) ----
  data.filename <- paste0(path, "/", datafile)

  # Patients with dose reduction (from Excel file)
  treatment.desc <- read_xlsx(
    data.filename,
    sheet = "Treatment",
    range = cell_cols("A:G"),
    col_types = c("text", "text", "numeric", "date", "date", "date", "numeric")
  )

  treatment.desc <- treatment.desc |>
    mutate(
      Dose = sprintf("%0.1f", Dose),
      TreatDose = sprintf("%0.1f", TreatDose),
      C1D1 = as_date(C1D1),
      StartDate = as_date(StartDate),
      EndDate = as_date(EndDate)
    ) |>
    # Transfer Dose from chr to factor
    mutate(
      Dose = factor(Dose, levels = DOSE.LEVELS),
      TreatDose = factor(TreatDose, levels = DOSE.LEVELS)
    )

  # Read patient information from Excel file
  # Patients without dose de-escalation (from patient.info)
  treatment.nodesc <- patient.info |>
    filter(is.na(isDoseDesc)) |>
    select(Patient, CurrentState, Dose, C1D1, EndDate) |>
    mutate(StartDate = C1D1, TreatDose = Dose)

  # treatment is used to draw the duration of treatment for each patient
  treatment <- bind_rows(treatment.desc, treatment.nodesc)

  # 计算研究治疗的持续时间
  ##  从C1D1到EndDate的持续时间，或者从C1D1到当前日期
  treatment <- treatment |>
    as.data.frame() |>
    mutate(
      Duration = ifelse(
        is.na(EndDate),
        as.numeric(difftime(today() + 1, C1D1, units = "weeks")),
        as.numeric(difftime(EndDate + 1, C1D1, units = "weeks"))
      ),
      Cont = ifelse(is.na(EndDate), 1, NA) # Arrow数据的Cont
    )

  # 从治疗数据中提取每位患者的首次给药信息
  firstdosing <- treatment |>
    arrange(Patient, desc(Duration)) |> # 首先按照患者和持续时间降序排列数据
    distinct(Patient, .keep_all = TRUE) |> # 然后去重，保留每位患者的第一条记录
    # 最后选择所需的列，包括患者、当前状态、剂量、C1D1和持续时间
    select(Patient, CurrentState, Dose, C1D1, Duration)

  # Keep only columns usded in the plot
  treatment <- treatment |>
    select(-CurrentState, -C1D1, -StartDate, -EndDate)

  # POINTS data (response/RECIST) ----
  # ## recist means response results
  recist <- read_targetlesion_info(path, datafile)

  recist <- recist |>
    as.data.frame() |>
    filter(No.eva > 0) |> # Remove Baseline
    select(-No.eva, -Sum) # Only keep RECIST information

  # response will be used to plot
  response <- left_join(recist, firstdosing, by = "Patient")

  response <- response |>
    mutate(EvaTime = as.numeric(difftime(Date.eva + 1, C1D1, units = "weeks")))

  # Date check ----
  # PD仍在治疗中，应该是错误的
  errors <- response |>
    filter(RECIST == "PD") |>
    filter(CurrentState == "治疗")
  if (nrow(errors) > 0) {
    print(errors)
    stop("PD patients with treatment!")
  }

  # Data Set ----
  treatment <- filter(treatment, Patient %in% plot.set$Patient)
  response <- filter(response, Patient %in% treatment$Patient)

  # Delete patients without duration data
  treatment <- treatment |>
    filter(!is.na(Duration))

  # max.duration is used to scale the X axis
  max.duration <- ceiling(max(
    max(treatment$Duration, na.rm = TRUE),
    max(response$EvaTime, na.rm = TRUE)
  ))

  # Plotting ----
  if (nrow(treatment) == 0) {
    print("No data for Swimmer plot!")
    return(NULL)
  }

  swimplot <- swimmer_plot(
    df = treatment,
    id = "Patient",
    end = "Duration",
    name_fill = "TreatDose",
    # id_order = "Dose", # Grouped by Dose
    # stratify = "Dose",
    alpha = 0.75,
    width = 0.8
  ) +
    theme_bw() +
    scale_fill_manual(values = DOSE.COLORS)

  if (nrow(treatment) > 120) {
    # 如果受试者太多，隐藏Y轴标签和刻度
    swimplot <- swimplot +
      theme(axis.text.y = element_blank())
  } else {
    # 如果显示所有受试者，要缩小ID的字体
    swimplot <- swimplot +
      theme(
        axis.text.y = element_text(size = 5)
        # axis.text.y = element_blank(),
        # axis.ticks.y = element_blank()
      )
  }

  swimplot_points <- swimplot +
    swimmer_points(
      df_points = response,
      id = "Patient",
      time = "EvaTime",
      name_shape = "RECIST",
      show.legend = c(fill = FALSE), # Hide fill legend
      # fill = "#C62128" # Fill color for CR
      # name_fill = "RECIST",
      # name_col = "RECIST",
      size = 1.5,
    ) +
    scale_shape_manual(
      name = "RECIST",
      values = RESP.SHAPES, # Shapes for CR, PR, SD, PD, NE
      breaks = RESP.LEVELS
    ) +
    guides(
      fill = guide_legend(title = "Dose"),
      shape = guide_legend(title = "Response")
    )

  swimplot_points_arrows <- swimplot_points +
    swimmer_arrows(
      df_arrows = treatment,
      id = "Patient",
      arrow_start = "Duration",
      cont = "Cont",
      name_col = "Cont",
      arrow_positions = c(0.3, 0.95), # the length of arrow
      type = "open",
      angle = 30,
      length = 0.05, # the length of arrow head
      # cex = 0.2,
      color = "black",
      show.legend = FALSE
    ) +
    guides(colour = guide_none()) +
    scale_color_discrete(drop = FALSE)

  swplot <- swimplot_points_arrows +
    labs(title = title, x = "", y = "Weeks") +
    geom_rug(show.legend = FALSE) +
    # ylim(0, max.duration) +
    scale_y_continuous(
      breaks = seq(0, max.duration + 1, 6),
      minor_breaks = seq(0, max.duration + 1, 3),
      expand = expansion(mult = c(0, 0.1)) # Remove the blanket between Y axis and bars
    ) +
    scale_x_discrete(expand = expansion(add = c(1, 1)))

  # # Move legend to the right bottom of the figure
  # swplot <- swplot +
  #   theme(legend.position = c(0.95, 0.05),
  #         legend.justification = c(1, 0),
  #         legend.background = element_rect(fill = "transparent"))

  # Print summary
  print("Swimmer plot")
  print(paste0(treatment |> summarise(np = n_distinct(Patient)), " patients"))
  print(paste0(treatment |> summarise(n = n()), " data points."))

  return(swplot)
}


#' Swimmer plot without dose reduction
#' @param plot.set: a data frame of patients information
#' @param title: plot title
#' @param path: the directory of the data file
#' @param datafile: the data file name
#' @return: an object of swimmer plot
ggswim2 <- function(
  plot.set,
  title = "Swimmer Plot",
  path = "D:/Work/TF/Ph1",
  datafile = "XNW28012_SubjectList.xlsx"
) {
  # Raw data (treatment) ----
  ## treatment is used to draw the duration of treatment for each patient
  data.filename <- paste0(path, datafile)

  # 计算日期之间的持续时间
  treatment <- patient.info |>
    as.data.frame() |>
    mutate(
      Duration = ifelse(
        is.na(EndDate),
        as.numeric(difftime(today() + 1, C1D1, units = "weeks")),
        as.numeric(difftime(EndDate + 1, C1D1, units = "weeks"))
      ),
      Cont = ifelse(is.na(EndDate), 1, NA) # Arrow数据的Cont
    )

  # 该代码用于从治疗数据中提取每位患者的首次给药信息
  # 首先按照患者和持续时间降序排列数据
  # 然后去重，保留每位患者的第一条记录
  # 最后选择所需的列，包括患者、当前状态、剂量、C1D1和持续时间
  firstdosing <- treatment |>
    arrange(Patient, desc(Duration)) |>
    distinct(Patient, .keep_all = TRUE) |>
    select(Patient, CurrentState, Dose, C1D1, Duration)

  # 去除患者数据中指定的列
  treatment <- treatment |>
    select(-CurrentState, -C1D1, -EndDate)

  # POINTS data (response/RECIST)
  recist <- read_targetlesion_info(path, datafile)

  recist <- recist |>
    as.data.frame() |>
    filter(No.eva > 0) |> # Remove Baseline
    select(-No.eva, -Sum) # Only keep RECIST info

  # response will be used to plot
  response <- left_join(recist, firstdosing, by = "Patient")

  response <- response |>
    mutate(EvaTime = as.numeric(difftime(Date.eva + 1, C1D1, units = "weeks")))

  # max.duration is used to scale the X axis
  max.duration <- max(treatment$Duration, na.rm = TRUE)

  # Data Set ----
  treatment <- filter(treatment, Patient %in% plot.set$Patient)
  response <- filter(response, Patient %in% treatment$Patient)

  # Del patients without duration data
  treatment <- treatment |> filter(!is.na(Duration))

  treatment$Dose <- as.character(treatment$Dose)

  # Plotting ----
  swimplot <- swimmer_plot(
    df = treatment,
    id = "Patient",
    end = "Duration",
    name_fill = "Dose",
    # id_order = "Dose", # Grouped by Dose
    # stratify = "Dose",
    alpha = 0.75,
    width = 0.8
  ) +
    theme_bw() +
    scale_fill_manual(values = DOSE.COLORS)
  # theme(axis.text.y = element_text(size = 5)
  # hide Y axis labels and ticks
  # axis.text.y = element_blank(),
  # axis.ticks.y = element_blank()
  # )

  swimplot_points <- swimplot +
    swimmer_points(
      df_points = response,
      id = "Patient",
      time = "EvaTime",
      name_shape = "RECIST",
      show.legend = c(fill = FALSE)
      # name_fill = "RECIST",
      # name_col = "RECIST",
      # size = 1.5
    ) +
    scale_shape_manual(
      name = "RECIST",
      values = RESP.SHAPES, # Shapes for CR, PR, SD, PD, NE
      breaks = RESP.LEVELS
    ) +
    guides(
      fill = guide_legend(title = "Dose"),
      shape = guide_legend(title = "Response")
    )

  swimplot_points_arrows <- swimplot_points +
    swimmer_arrows(
      df_arrows = treatment,
      id = "Patient",
      arrow_start = "Duration",
      cont = "Cont",
      name_col = "Cont",
      arrow_positions = c(0.3, 0.95), # the length of arrow
      type = "open",
      angle = 30,
      length = 0.05, # the length of arrow head
      # cex = 0.2,
      color = "black",
      show.legend = FALSE
    ) +
    guides(colour = guide_none()) +
    scale_color_discrete(drop = FALSE)

  swplot <- swimplot_points_arrows +
    labs(x = "", y = "Weeks") +
    geom_rug(show.legend = FALSE) +
    scale_y_continuous(
      breaks = seq(0, max.duration + 1, 6),
      minor_breaks = seq(0, max.duration + 1, 3)
    )

  # # Move legend to the right bottom of the figure
  # swplot <- swplot +
  #   theme(legend.position = c(0.95, 0.05),
  #         legend.justification = c(1, 0),
  #         legend.background = element_rect(fill = "transparent"))

  # Print summary
  print("Swimmer plot")
  print(paste0(treatment |> summarise(np = n_distinct(Patient)), " patients"))
  print(paste0(response |> summarise(n = n()), " response data points."))

  return(swplot)
}
